import os
import sys
import pandas as pd
import numpy as np
from tqdm import tqdm, trange
from matplotlib import pyplot as plt
import seaborn as sns
import json
import pathlib
from pathlib import Path

import h5py
import torch
import torch.nn as nn
import torch.nn.functional as F

from . import vgg16
from . import netvlad_v2
from .netvlad_v2 import EmbedNet
from ..configs import read_config_file
from ..utils import download_file

from ..configs import configs_path, root_path


import logging

logging.basicConfig(
    # filename="access.log",
    format="%(asctime)s - %(name)s - %(levelname)s -%(module)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S %p",
    level=logging.INFO,
)


def get_model():
    logging.info("building the netvlad model")

    if not os.path.exists(os.path.join(configs_path, "global.yml")):
        download_file(
            "https://gitee.com/GardenLu/garden_ibl/raw/master/configs/global.yml",
            os.path.join(configs_path, "global.yml"),
        )
    if not os.path.exists(os.path.join(configs_path, "netvlad.yml")):
        download_file(
            "https://gitee.com/GardenLu/garden_ibl/raw/master/configs/netvlad.yml",
            os.path.join(configs_path, "netvlad.yml"),
        )

    global_config = read_config_file(os.path.join(configs_path, "global.yml"))
    vlad_config = read_config_file(os.path.join(configs_path, "netvlad.yml"))

    matconvnet_path = os.path.join(
        global_config["root_path"],
        vlad_config["ckpt_path"],
        "vd16_offtheshelf_conv5_3_max.pth",
    )

    if not os.path.exists(os.path.join(global_config["root_path"], vlad_config["ckpt_path"])):
        os.makedirs(global_config["root_path"], vlad_config["ckpt_path"])

    if not os.path.exists(matconvnet_path):
        download_file(vlad_config["matconvnet"], matconvnet_path)

    initchache_path = os.path.join(
        global_config["root_path"],
        vlad_config["ckpt_path"],
        "vgg16_pitts_64_desc_cen.hdf5",
    )

    if not os.path.exists(global_config["root_path"], vlad_config["ckpt_path"]):
        os.makedirs(global_config["root_path"], vlad_config["ckpt_path"])

    if not os.path.exists(initchache_path):
        download_file(vlad_config["init_cache"], initchache_path)

    base_model = vgg16.vgg16(matconvnet=matconvnet_path)
    pool_layer = netvlad_v2.NetVLAD(dim=base_model.feature_dim)

    initchache = initchache_path

    # pool_layer.clsts =
    with h5py.File(initchache, mode="r") as h5:
        pool_layer.clsts = h5.get("centroids")[...]
        pool_layer.traindescs = h5.get("descriptors")[...]
        pool_layer._init_params()

    model = EmbedNet(base_model, pool_layer)

    model = model.cuda()
    # model = nn.parallel.DistributedDataParallel(
    #     model, device_ids=[0], output_device=0, find_unused_parameters=True
    # )

    return model
